import json
import matplotlib.pyplot as plt
import os
import argparse

# Read and parse the JSON file
def read_json_file(file_path):
    with open(file_path, 'r') as file:
        data = [json.loads(line.strip()) for line in file if line.strip()]
    return data

# Collect D2C delay data
def collect_d2c_latency_data(data):
    d2c_data = []
    for item in data:
        event = item['event']
        if event == 'C' and item['blk_dts'] != 0 and item['blk_cts'] != 0:
            latency = item['blk_cts'] - item['blk_dts']
            d2c_data.append((item['ts'], latency))
    return d2c_data

# Collect G2I delay data
def collect_g2i_latency_data(data):
    g2i_data = []
    for item in data:
        event = item['event']
        if event == 'I' and item['blk_gts'] != 0 and item['blk_its'] != 0:
            latency = item['blk_its'] - item['blk_gts']
            g2i_data.append((item['ts'], latency))
    return g2i_data


# Collect I2D delay data
def collect_i2d_latency_data(data):
    i2d_data = []
    for item in data:
        event = item['event']
        if event == 'D' and item['blk_its'] != 0 and item['blk_dts'] != 0:
            latency = item['blk_dts'] - item['blk_its']
            i2d_data.append((item['ts'], latency))
    return i2d_data

# Plot a scatter plot and save it
def plot_scatter(data, output_file, plot_title, ylabel):
    times, latencies = zip(*data)  
    plt.figure(figsize=(10, 6))  
    plt.scatter(times, latencies, marker='o')  
    plt.xlabel('Time (nanosecond)')
    plt.ylabel(ylabel)
    plt.title(plot_title)
    plt.grid(True)
    plt.savefig(output_file)  
    plt.close()  


def main():
    parser = argparse.ArgumentParser(description='Analyze and plot event data latencies.')
    parser.add_argument('-t', '--type', type=str, choices=['I2D', 'D2C', 'G2I'],
                        help='Specify the type of latency to analyze (I2G, D2C, G2I).')
    parser.add_argument('-p', '--plot', type=str, default='latency_scatter.png',
                        help='Specify the output file name for the plot image.')
    parser.add_argument('-i', '--input', type=str,
                    help='Specify the input JSON file path.')
    args = parser.parse_args()

    # Read data
    data = read_json_file(args.input)

    # Collect data based on specified types
    if args.type == 'D2C':
        latency_data = collect_d2c_latency_data(data)
        plot_title = 'D2C Latency over Time'
        ylabel = 'D2C Latency (millisecond)'
    elif args.type == 'G2I':
        latency_data = collect_g2i_latency_data(data)
        plot_title = 'G2I Latency over Time'
        ylabel = 'G2I Latency (millisecond)'
    elif args.type == 'I2D':
        latency_data = collect_i2d_latency_data(data)
        plot_title = 'I2D Latency over Time'
        ylabel = 'I2D Latency (millisecond)'

    output_dir = 'output_plots'
    os.makedirs(output_dir, exist_ok=True)

    output_file = os.path.join(output_dir, args.plot)

    # Plot and save the chart
    plot_scatter(latency_data, output_file, plot_title, ylabel)

if __name__ == '__main__':
    main()